#!/usr/bin/env python3
# policy/pursuit_policy.py
# 策略聚合类：组合信念、计算、远程服务器
import os
import numpy as np
import networkx as nx

from .belief_manager import BeliefManager
from .strategy_calculator import StrategyCalculator
from .remote_policy_server import RemotePolicyServer
from map_manager.map_manage import MapManager

import logging

logger = logging.getLogger(__name__)


class PursuitPolicy:
    def __init__(self, map_manager: MapManager, logger=None):
        if logger is None:
            logger = logging.getLogger(__name__)
        self.logger = logger
        self.map_manager = map_manager
        self.load_graph_data_from_manager()
        self.capture_count = 0
        # 初始化子组件
        self.belief_manager = BeliefManager(
            self.node_num, self.adj_matrix, obs_dist=4, pursuer_mode=1
        )
        self.strategy_calc = StrategyCalculator(
            self.adj_matrix, self.f_matrix, self.dist_3d, self.dbin_available
        )
        self.remote_server = RemotePolicyServer(self.adj_matrix, self.node_num, logger)
        self.logger.info("PursuitPolicy 初始化完成")

    def load_graph_data_from_manager(self):
        """从MapManager加载数据"""
        self.node_num = self.map_manager.node_num
        self.node_coordinates = self.map_manager.node_coordinates
        self.adj_matrix = self.map_manager.adj_matrix.copy()
        logger.info(
            f"Load graph from MapManager: node_num={self.node_num}, node_coordinates={len(self.node_coordinates)}, adj_matrix={self.adj_matrix.shape}"
        )
        # 加载距离矩阵
        self.map_dir_path = os.path.join(os.path.dirname(__file__), "../datas/")
        dist_npy_path = os.path.join(self.map_dir_path, "D_bin.npy")
        try:
            if os.path.exists(dist_npy_path):
                self.dist_3d = np.load(dist_npy_path)
                if self.dist_3d.shape != (
                    self.node_num + 1,
                    self.node_num + 1,
                    self.node_num + 1,
                ):
                    raise ValueError(
                        f"距离矩阵形状不匹配, 期望 {(self.node_num + 1, self.node_num + 1, self.node_num + 1)}, 实际 {self.dist_3d.shape}"
                    )
                self.dbin_available = True
                self.logger.info("已加载距离张量")
                # 计算最短路径
                self.compute_shortest_paths()
            else:
                self.dbin_available = False
                self.dist_3d = None
                self.logger.warning(
                    "距离矩阵文件不存在，禁用距离张量功能. 路径: " + dist_npy_path
                )
        except Exception as e:
            self.logger.warning(f"加载距离矩阵失败: {e}")
            self.dbin_available = False
            self.dist_3d = None
        # 更新子组件
        self.belief_manager = BeliefManager(self.node_num, self.adj_matrix)
        self.strategy_calc = StrategyCalculator(
            self.adj_matrix, self.f_matrix, self.dist_3d, self.dbin_available
        )

    def compute_shortest_paths(self):
        """计算Floyd-Warshall矩阵"""
        G = nx.from_numpy_array(self.adj_matrix, create_using=nx.DiGraph())
        dist_matrix = nx.floyd_warshall_numpy(G)
        INF = 255
        dist_matrix[np.isinf(dist_matrix)] = INF
        self.f_matrix = np.rint(dist_matrix).astype(int)
        np.fill_diagonal(self.f_matrix, 0)
        self.logger.info("最短路径矩阵计算完成")

    def reset(self):
        """重置策略"""
        self.logger.info("PursuitPolicy 重置中...")
        self.load_graph_data_from_manager()
        self.capture_count = 0
        self.belief_manager.belief_pr.fill(0.0)
        self.belief_manager.belief_pos.fill(False)
        # 如何远程服务器已经初始化成功，则发送图结构更新到远程服务器
        if self.remote_server.policy_socket is None:
            return  # 未连接则跳过
        self.remote_server.adj_matrix = self.adj_matrix  # 更新服务器内部矩阵
        self.remote_server.node_num = self.node_num
        self.remote_server.send_graph_update()
        self.logger.info("PursuitPolicy 重置完成")

    def init_belief(self, evader_node):
        self.belief_manager.init_belief(evader_node)

    def update_belief(self, evader_node, p1_node, p2_node, visible_after_move=False):
        self.belief_manager.update_belief(
            evader_node, p1_node, p2_node, self.f_matrix, visible_after_move
        )

    def calculate_evader_move(self, evader_node, p1_node, p2_node):
        """独立计算逃避者移动（本地运算）"""
        return self.strategy_calc.calculate_evader_move(evader_node, p1_node, p2_node)

    def calculate_pursuers_move(self, evader_node, p1_node, p2_node, belief_manager):
        """独立计算追捕者移动（包括信念更新、捕获检查、远程计算）"""
        visible_after = belief_manager.visible(
            p1_node, p2_node, evader_node, self.f_matrix
        )
        belief_manager.update_belief(
            evader_node, p1_node, p2_node, self.f_matrix, visible_after
        )

        if self.strategy_calc.is_captured(evader_node, p1_node, p2_node):
            return p1_node, p2_node  # 已捕获，无需移动

        belief_captured = any(
            belief_manager.belief_pos[l]
            and self.strategy_calc.cond1(p1_node, p2_node, l)
            for l in range(self.node_num)
        )
        if belief_captured:
            return p1_node, p2_node

        return self.strategy_calc._remote_calculate_pursuers_move(
            evader_node, p1_node, p2_node, self.remote_server
        )

    def is_captured(self, evader_node, p1_node, p2_node):
        return self.strategy_calc.is_captured(evader_node, p1_node, p2_node)

    def visible(self, p1, p2, l):
        return self.strategy_calc.visible(p1, p2, l)

    def get_belief_pr(self):
        return self.belief_manager.get_belief_pr()

    def get_belief_pos(self):
        return self.belief_manager.get_belief_pos()

    def get_capture_count(self):
        return self.capture_count

    def increment_capture_count(self):
        self.capture_count += 1

    # 代理MapManager方法
    def position_to_node(self, x, y):
        return self.map_manager.position_to_node(x, y)

    def node_to_position(self, node_id):
        return self.map_manager.node_to_position(node_id)
